import os
os.environ['CUDA_VISIBLE_DEVICES']='0'
import util_for_huggingface
import torch,pdb
# from diffusers import  FluxPipeline
from diffusers import FluxFillPipeline, FluxPriorReduxPipeline,ReduxImageEncoder
from diffusers.utils import load_image
from util_flux import horizontal_concat_images,process_img_1024
from PIL import Image
# from image_gen_aux import DepthPreprocessor

'''
20250714
 测试 train_flux_depth_local2.py 的训练结果
 所需数据 :
    oriimg
    localmask
    localimg
 base-model: FILL
 所需处理
    localimg -> lotus = dpth
    siglip(dpth) siglip(localimg)  => prompt_emb
    oriimg + localmask             => for concat img2 as input
    localmask + black              => mask
'''

FLUX_FILL='/home/shengjie/ckp/FLUX.1-Fill-dev'
DEPTH_PREDCITION='/home/shengjie/ckp/depth-anything-large-hf'
FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX_DEPTH='/home/shengjie/ckp/FLUX.1-Depth-dev'
FLUX_ADAPTER = '/home/shengjie/ckp/flux-ip-adapter-v2'
FLUX_ADAPTER_ENCODER = '/home/shengjie/ckp/clip-vit-large-patch14'
FLUX='/data/models/FLUX___1-dev'

##############  clip 获得 propmt_embeds ##########
from util_for_clip import get_clip_model
from util_for_os import ose,osj

## clip for prompt_embeds
clip_processor,clip_model,maxtokens = get_clip_model()

# examples_dir = '/data/shengjie/style_zhenzhi/'
# save_dir = '/data/shengjie/synthesis_zhenzhi/'
types = ['collar','sleeve','pockets']
choose_type = types[0]
data_root = '/mnt/nas/shengjie2/datasets'
get_ori_data_dir = lambda t : osj( data_root , f'cloth_{t}_balanced' )
get_localimg_data_dir = lambda t : osj( data_root , f'cloth_{t}_localimg' )
get_localimg_mask_data_dir = lambda t : osj( data_root , f'cloth_{t}_localimg_mask' )
get_localimg_canny_data_dir = lambda t : osj( data_root , f'cloth_{t}_localimg_canny' )

ori_data_dir = get_ori_data_dir(choose_type)
localimg_data_dir = get_localimg_data_dir(choose_type)
localimg_mask_data_dir = get_localimg_mask_data_dir(choose_type)
localimg_canny_data_dir = get_localimg_canny_data_dir(choose_type)


dtype = torch.bfloat16
device = 'cuda'

ckp_id = 5000
ckp_dir = 'depth_local_output_20250714'
get_local_embeder_path = lambda ckp_id : f'/mnt/nas/shengjie/{ckp_dir}/checkpoint-{ckp_id}/image_encoder_1.bin'
get_dpth_embeder_path = lambda ckp_id : f'/mnt/nas/shengjie/{ckp_dir}/checkpoint-{ckp_id}/image_encoder_2.bin'
get_lora_path = lambda ckp_id : f'/mnt/nas/shengjie/{ckp_dir}/checkpoint-{ckp_id}/pytorch_lora_weights.safetensors'
local_embder_ckp_path = get_local_embeder_path(ckp_id)
dpth_embder_ckp_path = get_dpth_embeder_path(ckp_id)
lora_path = get_lora_path(ckp_id)

from lotus.app_infer_depth import load_pipe_d,get_depth_by_lotus_d
task_name = 'n'
depth_pipe = load_pipe_d(task_name).to(device,dtype=torch.bfloat16)

pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                        FLUX_REDUX, 
                                        torch_dtype=torch.bfloat16).to(device)
local_embder = ReduxImageEncoder().to(device,dtype=dtype)
dpth_embder = ReduxImageEncoder().to(device,dtype=dtype)

pipe = FluxFillPipeline.from_pretrained(FLUX_FILL, 
                                        torch_dtype=torch.bfloat16).to("cuda")

local_embder_ckp = torch.load(local_embder_ckp_path)
local_embder.load_state_dict(local_embder_ckp)
dpth_embder_ckp = torch.load(dpth_embder_ckp_path)
dpth_embder.load_state_dict(dpth_embder_ckp)
pipe.load_lora_weights(lora_path)

# load lora

def get_masked_img(img,mask):
    mask = mask.convert('L')
    background = Image.new("RGB", img.size, 'white')
    from PIL import ImageOps
    masked_img = Image.composite(img , background , ImageOps.invert( mask ) )
    return masked_img
def horizontal_concat_images(images,):
    # 计算总宽度 (所有图片宽度之和 + 间距)
    widths, heights = zip(*(img.size for img in images))
    total_width = sum(widths)
    max_height = max(heights)
    
    # 创建新图片 (白色背景)
    new_img = Image.new('RGB', (total_width, max_height), color='white')
    
    # 粘贴图片并绘制分隔线
    x_offset = 0
    for i, img in enumerate(images):
        new_img.paste(img, (x_offset, 0))
        
        x_offset += img.width
    
    return new_img
count = 0
prename = None
for entry in os.scandir(ori_data_dir):
    if not entry.is_file():continue
    filename = entry.name
    if not filename.endswith('.jpg'):continue

    count+=1
    print('\rprocess idx: ',count,end='',flush=True)

    ori_filepath = osj( ori_data_dir , filename )
    if prename is None:
        prename = filename
        continue
    # local img 1 for depth
    # local img 2 for redux
    localimg_filepath = osj( localimg_data_dir , prename )
    localimg_filepath2 = osj( localimg_data_dir , filename )
    localimg_mask_filepath = osj( localimg_mask_data_dir , filename.replace('.jpg','.png') )
    localimg_canny_filepath = osj( localimg_canny_data_dir , filename )
    # localimg_normal_filepath = osj( localimg_normal_data_dir , filename )

    prename = filename

    assert ose(ori_filepath),ori_filepath
    assert ose(localimg_filepath),localimg_filepath
    assert ose(localimg_filepath2),localimg_filepath2 # for redux
    assert ose(localimg_mask_filepath),localimg_mask_filepath
    assert ose(localimg_canny_filepath),localimg_canny_filepath
    # assert ose(localimg_normal_filepath),localimg_normal_filepath

    ori_imgae = process_img_1024(ori_filepath).convert("RGB")
    local_img = process_img_1024(localimg_filepath).convert("RGB")
    local_img2 = process_img_1024(localimg_filepath2).convert("RGB")
    local_mask = process_img_1024(localimg_mask_filepath,pad_color=(0,0,0)).convert("RGB")

    # depth
    local_depth_image = get_depth_by_lotus_d(depth_pipe,
                                            local_img,
                                            None,device,)
                
    # prompt emb
    local_image_latents = pipe_prior_redux.encode_image( local_img2 , pipe_prior_redux.device , 1 )
    local_image_latents = local_image_latents.to(device,dtype = dtype)   
    local_depth_latents = pipe_prior_redux.encode_image( local_depth_image , pipe_prior_redux.device , 1  )
    local_depth_latents = local_depth_latents.to(device,dtype = dtype)   
    prompt_embeds2 = 0.5*local_embder(local_image_latents)['image_embeds']+\
                    0.5*dpth_embder(local_depth_latents)['image_embeds']
    prompt_embeds1 = torch.zeros(1,512,4096)\
                    .to(device,dtype = dtype)
    prompt_embeds = torch.concat( [prompt_embeds1,prompt_embeds2],dim=1 )\
                    .to(device,dtype = dtype)  
    pooled_prompt_embeds = torch.zeros(1,768)\
                    .to(device,dtype = dtype)
    text_ids = torch.zeros(prompt_embeds.shape[1], 3)\
                    .to(device,dtype = dtype)


    # for input
    local_masked_img = get_masked_img(ori_imgae,local_mask)
    black_image = Image.new("RGB", local_masked_img.size , (0, 0 , 0) )
    input_img = horizontal_concat_images([local_masked_img,local_masked_img])
    # for mask
    black_image = Image.new("RGB", input_img.size , (0, 0, 0))
    black_image.paste(local_mask, (0, 0))
    mask = black_image


    pdb.set_trace()
    # infer
    with torch.no_grad():
        image = pipe(
            image=input_img,
            mask_image=mask,
            height=input_img.size[1],
            width=input_img.size[0],

            prompt_embeds=prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            
            guidance_scale=1.0,
            num_inference_steps = 10,
            # depth_condition=packed_depth_latents,
            # texture_condition=packed_texture_latents,
        ).images[0]


    # image.save('tmp.jpg')
    width,height = image.size
    image_concat = Image.new('RGB' , (width,height*2) , (255,255,255))
    image_concat.paste(image , (0,0))
    image_concat.paste(local_img , (0,height))
    
    image_concat.save('tmp_test_local.jpg')
    pdb.set_trace()

    # break


